/**
 * 内部头文件，不对用户提供
 */

#ifndef __WZZ_PICPROC_INTERNAL_HPP__
#define __WZZ_PICPROC_INTERNAL_HPP__

#include <vector>
#include <unordered_map>
#include <string>

#include <onnxruntime_cxx_api.h>


namespace wzz
{

    // 定义一个string_map的哈希表类型用于存储键为string，值为T类型的shared_ptr
    template<typename T>
    using string_map =  std::unordered_map<std::string, std::shared_ptr<T>>; 

    // 对于每一个类型T都只维护一个该全局类型的映射表
    template <typename T>
    string_map<T> & get_global_map() noexcept {
        static wzz::string_map<T> global_map;
        return global_map;
    }

    // 根据键获取对应类型的指针，如果该键值没有对应的值则返回nullptr
    template<typename T>
    T * get_map_value(const std::string & key) noexcept {
        if (get_global_map<T>().count(key) == 0) {
            return nullptr;
        }

        return get_global_map<T>()[key].get();
    }

    // 向全局表中注册一个T类型的值
    template<typename T, typename... Args>
    void register_map_value(std::string key, Args&& ... args) {
        if (get_global_map<T>().count(key) == 1) {
            std::string err_string;
            err_string += "register key:" + key + " alread registered";
            throw std::runtime_error(err_string);
        }
        get_global_map<T>()[key] = std::make_shared<T>(std::forward(args)...);
    }

    // onnx模型
    class Model {
    public:
        std::tuple<std::vector<float>, std::vector<int64_t>> 
        operator()(
            const std::vector<float> &input, const std::vector<int64_t> &input_shape
        ){
            Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
                m_memory_info, m_input.data(), m_input.size(),
                m_input_shape.data(), m_input_shape.size()
            );

            Ort::Value output_tensor = Ort::Value::CreateTensor<float> (
                m_memory_info, m_output.data(), m_output.size(),
                m_output_shape.data(), m_output_shape.size()
            );

            m_session.Run(
                Ort::RunOptions{nullptr}, 
                m_input_names, &input_tensor, m_input_count,
                m_output_names, &output_tensor, m_output_count
            );

            return std::make_tuple(m_output, m_input_shape);
        }
        
    private:
        std::vector<int64_t> m_input_shape;
        std::vector<int64_t> m_output_shape;
        std::vector<float> m_input;
        std::vector<float> m_output;

        Ort::Env m_env;
        Ort::Session m_session;
        Ort::MemoryInfo m_memory_info;

        const char **m_input_names;
        const char **m_output_names;
        int m_input_count;
        int m_output_count;
    };
    
} // namespace wzz


#endif